import yaml
from pathlib import Path
import subprocess

def load_config():
    with open("train_config.yaml") as f:
        return yaml.safe_load(f)

def build_command(config):
    base = config["base"]
    train = config["training"]
    
    run_name = f"{base['name']}_e{base['epoch']}_l{base['learning_rate']}"
    output_dir = f"./parameters/{run_name}"
    
    return [
        "accelerate", "launch",
        "--config_file", "config.yaml",
        "train.py",
        "--run_name", run_name,
        "--pretrained_model_name_or_path", base["model_path"],
        "--train_dataset_path", f"{base['data_path']}/{base['name']}.parquet",
        "--eval_dataset_path", f"{base['data_path']}/{base['name']}.parquet",
        "--num_train_epochs", str(base["epoch"]),
        "--learning_rate", str(base["learning_rate"]),
        "--output_dir", output_dir,
        "--eval_steps", str(train["eval_steps"]),
        "--save_steps", str(train["save_steps"]),
        "--max_length", str(train["max_length"]),
        "--per_device_train_batch_size", str(train["batch_size"]),
        "--per_device_eval_batch_size", str(train["batch_size"])
    ]

def main():
    config = load_config()
    cmd = build_command(config)
    subprocess.run(cmd)

if __name__ == "__main__":
    main()